#include "hnurm_armor/Processor.h"
#include "hnurm_armor/DataType.hpp"
#include <angles/angles.h>

#include <opencv2/core/eigen.hpp>

namespace hnurm
{
    void Processor::ProcessArmor(
        std::vector<Armor> &armors_msg, hnurm_interfaces::msg::VisionRecvData &recv_data, TargetInfo &target_msg
    )
    {
        // Transform armor position from image frame to world coordinate
        for (auto &armor: armors_msg)
        {
            tsolver->GetTransformation(
                armor, (float) angles::from_degrees(recv_data.pitch), (float) angles::from_degrees(recv_data.yaw)
            );
        }

        auto time = std::chrono::steady_clock::now();

        if (tracker->tracker_state == Tracker::LOST)
        {
            tracker->init(armors_msg);
        }
        else
        {
            dt_ = (double) std::chrono::duration_cast<Us>(time - time_elapsed_).count() / 1e6;
            tracker->update(armors_msg);
        }
        time_elapsed_ = time;
        const auto state = tracker->target_state;
        // clang-format off
        target_msg.x = (float) state(0);
        target_msg.vx = (float) state(1);
        target_msg.y = (float) state(2);
        target_msg.vy = (float) state(3);
        target_msg.z = (float) state(4);
        target_msg.vz = (float) state(5);
        target_msg.yaw = (float) state(6);
        target_msg.w_yaw = (float) state(7);
        target_msg.radius_1 = (float) state(8);
        target_msg.radius_2 = (float) tracker->another_r;
        target_msg.dz = (float) tracker->dz;
        // clang-format on
        if (!armors_msg.empty())
        {
            if (target_msg.w_yaw < 1)
            {
                auto gimbal_yaw = recv_data.yaw;
                std::sort(armors_msg.begin(), armors_msg.end(), [&](const Armor &armor1, const Armor &armor2)
                {
                    Eigen::Vector3d euler1, euler2;
                    this->tsolver->rotationMatrixToEulerAngles(armor1.rotation, euler1);
                    this->tsolver->rotationMatrixToEulerAngles(armor2.rotation, euler2);
                    return abs(euler1(1) - gimbal_yaw) < abs(euler2(1) - gimbal_yaw);
                });

                target_msg.x = armors_msg[0].position.x();
                target_msg.y = armors_msg[0].position.y();
                target_msg.z = armors_msg[0].position.z();
            }
        }
    }

    Processor::Processor(rclcpp::Node::SharedPtr node)
        : node_(node)
    {

        dt_ = 0.01;

        tracker = std::make_unique<Tracker>(node_);
        tsolver = std::make_unique<TSolver>(node_);

        // EKF
        // xa = x_armor, xc = x_robot_center
        // state: xc, v_xc, yc, v_yc, za, v_za, yaw, v_yaw, r
        // measurement: xa, ya, za, yaw
        // f - Process function
        auto f = [this](const Eigen::VectorXd &x)
        {
            Eigen::VectorXd x_new = x;
            x_new(0) += x(1) * dt_;
            x_new(2) += x(3) * dt_;
            x_new(4) += x(5) * dt_;
            x_new(6) += x(7) * dt_;
            return x_new;
        };

        // J_f - Jacobian of process function
        auto j_f = [this](const Eigen::VectorXd &)
        {
            Eigen::MatrixXd f(9, 9);
            // clang-format off
            f << 1, dt_, 0, 0, 0, 0, 0, 0, 0,
                0, 1, 0, 0, 0, 0, 0, 0, 0,
                0, 0, 1, dt_, 0, 0, 0, 0, 0,
                0, 0, 0, 1, 0, 0, 0, 0, 0,
                0, 0, 0, 0, 1, dt_, 0, 0, 0,
                0, 0, 0, 0, 0, 1, 0, 0, 0,
                0, 0, 0, 0, 0, 0, 1, dt_, 0,
                0, 0, 0, 0, 0, 0, 0, 1, 0,
                0, 0, 0, 0, 0, 0, 0, 0, 1;
            // clang-format on
            return f;
        };

        // h - Observation function
        auto h = [](const Eigen::VectorXd &x)
        {
            Eigen::VectorXd z(4);
            double xc = x(0), yc = x(2), yaw = x(6), r = x(8);
            z(0) = xc - r * cos(yaw);  // xa
            z(1) = yc - r * sin(yaw);  // ya
            z(2) = x(4);               // za
            z(3) = x(6);               // yaw
            return z;
        };
        // J_h - Jacobian of observation function
        auto j_h = [](const Eigen::VectorXd &x)
        {
            Eigen::MatrixXd h(4, 9);
            double yaw = x(6), r = x(8);
            // clang-format off
            //    xc   v_xc yc   v_yc za   v_za yaw            v_yaw r
            h << 1, 0, 0, 0, 0, 0, r * sin(yaw), 0, -cos(yaw),
                0, 0, 1, 0, 0, 0, -r * cos(yaw), 0, -sin(yaw),
                0, 0, 0, 0, 1, 0, 0, 0, 0,
                0, 0, 0, 0, 0, 0, 1, 0, 0;
            // clang-format on
            return h;
        };

        // tracker_config_node["s2_q_xyz"] >> s2qxyz_;
        // tracker_config_node["s2_q_yaw"] >> s2qyaw_;
        // tracker_config_node["s2_q_r"] >> s2qr_;
        s2qxyz_ = node->declare_parameter("processor.tracker.s2_q_xyz", 0.05);
        s2qyaw_ = node->declare_parameter("processor.tracker.s2_q_yaw", 5.0);
        s2qr_ = node->declare_parameter("processor.tracker.s2_q_r", 80.0);

        // update_Q - process noise covariance matrix
        auto u_q = [this]()
        {
            Eigen::MatrixXd q(9, 9);
            double t = dt_, x = s2qxyz_, y = s2qyaw_, r = s2qr_;
            double q_x_x = pow(t, 4) / 4 * x, q_x_vx = pow(t, 3) / 2 * x, q_vx_vx = pow(t, 2) * x;
            double q_y_y = pow(t, 4) / 4 * y, q_y_vy = pow(t, 3) / 2 * x, q_vy_vy = pow(t, 2) * y;
            double q_r = pow(t, 4) / 4 * r;
            // clang-format off
            //    xc      v_xc    yc      v_yc    za      v_za    yaw     v_yaw   r
            q << q_x_x, q_x_vx, 0, 0, 0, 0, 0, 0, 0,
                q_x_vx, q_vx_vx, 0, 0, 0, 0, 0, 0, 0,
                0, 0, q_x_x, q_x_vx, 0, 0, 0, 0, 0,
                0, 0, q_x_vx, q_vx_vx, 0, 0, 0, 0, 0,
                0, 0, 0, 0, q_x_x, q_x_vx, 0, 0, 0,
                0, 0, 0, 0, q_x_vx, q_vx_vx, 0, 0, 0,
                0, 0, 0, 0, 0, 0, q_y_y, q_y_vy, 0,
                0, 0, 0, 0, 0, 0, q_y_vy, q_vy_vy, 0,
                0, 0, 0, 0, 0, 0, 0, 0, q_r;
            // clang-format on
            return q;
        };

        r_xyz_factor = node->declare_parameter("processor.tracker.r_xyz_factor", 4e-4);
        r_yaw = node->declare_parameter("processor.tracker.r_yaw", 5e-3);

        // update_R - measurement noise covariance matrix
        auto u_r = [this](const Eigen::VectorXd &z)
        {
            Eigen::DiagonalMatrix<double, 4> r;
            double x = r_xyz_factor;
            r.diagonal() << abs(x * z[0]), abs(x * z[1]), abs(x * z[2]), r_yaw;
            return r;
        };

        // P - error estimate covariance matrix
        Eigen::DiagonalMatrix<double, 9> p0;
        p0.setIdentity();

        // init ekf
        tracker->ekf = ExtendedKalmanFilter{f, h, j_f, j_h, u_q, u_r, p0};

        tf_buffer_ = std::make_shared<tf2_ros::Buffer>(node_->get_clock());
        tf_listener_ = std::make_shared<tf2_ros::TransformListener>(*tf_buffer_);
    }

}  // namespace hnurm